# 股票API工具
import requests
import os
from datetime import datetime, timedelta

class StockAPI:
    def __init__(self):
        self.api_key = os.environ.get('STOCK_API_KEY')
        self.base_url = 'https://api.10jqka.com.cn/v6'
        
    def get_stock_basic_info(self, stock_code):
        """
        获取股票基本信息
        :param stock_code: 股票代码
        :return: 股票基本信息
        """
        url = f'{self.base_url}/stock/info'
        params = {
            'code': stock_code,
            'key': self.api_key
        }
        
        try:
            response = requests.get(url, params=params)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f'获取股票基本信息失败: {e}')
            return None
        
    def get_historical_data(self, stock_code, days=30):
        """
        获取股票历史数据
        :param stock_code: 股票代码
        :param days: 获取天数
        :return: 历史数据
        """
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')
        
        url = f'{self.base_url}/kline'
        params = {
            'code': stock_code,
            'begin': start_date,
            'end': end_date,
            'fields': 'date,open,close,high,low,volume,amount',
            'key': self.api_key
        }
        
        try:
            response = requests.get(url, params=params)
            response.raise_for_status()
            data = response.json()
            return data
        except Exception as e:
            print(f'获取股票历史数据失败: {e}')
            # 返回模拟数据，实际应用中应该处理错误
            return self._generate_sample_data(stock_code, days)
        
    def get_real_time_data(self, stock_code):
        """
        获取股票实时数据
        :param stock_code: 股票代码
        :return: 实时数据
        """
        url = f'{self.base_url}/realstock'
        params = {
            'code': stock_code,
            'key': self.api_key
        }
        
        try:
            response = requests.get(url, params=params)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f'获取股票实时数据失败: {e}')
            return None
        
    def _generate_sample_data(self, stock_code, days):
        """
        生成模拟股票数据
        :param stock_code: 股票代码
        :param days: 数据天数
        :return: 模拟数据
        """
        sample_data = {
            'code': stock_code,
            'name': 'Sample Stock',
            'klines': []
        }
        
        start_price = 100.0
        for i in range(days):
            date = (datetime.now() - timedelta(days=days-i)).strftime('%Y-%m-%d')
            open_price = round(start_price + (i % 5 - 2) * 0.5, 2)
            high_price = round(open_price + 1.0, 2)
            low_price = round(open_price - 1.0, 2)
            close_price = round(open_price + (i % 3 - 1) * 0.8, 2)
            volume = 1000000 + i * 10000
            amount = close_price * volume
            
            sample_data['klines'].append({
                'date': date,
                'open': open_price,
                'high': high_price,
                'low': low_price,
                'close': close_price,
                'volume': volume,
                'amount': amount
            })
            
            start_price = close_price
        
        return sample_data